"""YOLOv8 detector wrapper with parsing and class filtering."""
from __future__ import annotations
from dataclasses import dataclass
from typing import List, Tuple, Dict, Any

try:
    from ultralytics import YOLO
except Exception:  # ultralytics may be missing at bootstrap
    YOLO = None  # type: ignore


@dataclass
class Detection:
    cls: str
    conf: float
    bbox_xyxy: Tuple[int, int, int, int]

    @property
    def center(self) -> Tuple[int, int]:
        x1, y1, x2, y2 = self.bbox_xyxy
        return int((x1 + x2) / 2), int((y1 + y2) / 2)


class YoloDetector:
    def __init__(self, weights: str, device: str = "cpu"):
        if YOLO is None:
            raise RuntimeError("Ultralytics not installed. Please install ultralytics.")
        self.model = YOLO(weights)
        self.device = device
        # name index mapping
        self.names: Dict[int, str] = self.model.model.names if hasattr(self.model, "model") else self.model.names

    def _class_indices_by_names(self, wanted: List[str] | None) -> List[int] | None:
        if not wanted:
            return None
        inv = {v: k for k, v in self.names.items()}
        idxs = [inv[n] for n in wanted if n in inv]
        return idxs or None

    def predict(self, img, imgsz: int, conf: float, iou: float, class_names: List[str] | None = None):
        classes = self._class_indices_by_names(class_names)
        res = self.model.predict(img, imgsz=imgsz, conf=conf, iou=iou, device=self.device, classes=classes, verbose=False)
        return res

    def detect_persons_vehicles(self, img, imgsz: int, conf: float, iou: float,
                                 person_label: str = "person",
                                 vehicle_labels: List[str] | None = None) -> Tuple[List[Detection], List[Detection]]:
        if vehicle_labels is None:
            vehicle_labels = ["car", "bus", "truck", "motorcycle", "bicycle"]
        wanted = [person_label] + vehicle_labels
        res = self.predict(img, imgsz, conf, iou, class_names=wanted)
        if not res:
            return [], []
        r = res[0]
        boxes = getattr(r, "boxes", None)
        persons: List[Detection] = []
        vehicles: List[Detection] = []
        if boxes is None:
            return persons, vehicles
        import numpy as np  # local import
        xyxy = boxes.xyxy.cpu().numpy().astype(int)
        confs = boxes.conf.cpu().numpy()
        clsi = boxes.cls.cpu().numpy().astype(int)
        for (x1, y1, x2, y2), cf, ci in zip(xyxy, confs, clsi):
            name = self.names.get(int(ci), str(ci))
            det = Detection(cls=name, conf=float(cf), bbox_xyxy=(int(x1), int(y1), int(x2), int(y2)))
            if name == person_label:
                persons.append(det)
            elif name in vehicle_labels:
                vehicles.append(det)
        return persons, vehicles
